"""
指数K线数据模块
获取指数K线数据
"""
import time
from typing import Dict, List, Optional, Tuple
from datetime import datetime, timedelta
from loguru import logger

try:
    import okx.MarketData as MarketData
except ImportError:
    logger.warning("okx SDK未安装，请运行: pip install okx")
    MarketData = None


class IndexCandlesticksManager:
    """指数K线数据管理器"""
    
    # 限速配置：20次/2s（普通K线）
    RATE_LIMIT_REQUESTS = 20
    RATE_LIMIT_WINDOW = 2.0  # 秒
    
    # 历史K线限速配置：10次/2s
    HISTORY_RATE_LIMIT_REQUESTS = 10
    HISTORY_RATE_LIMIT_WINDOW = 2.0  # 秒
    
    # K线周期映射
    BAR_SIZES = {
        "1m": "1分钟",
        "3m": "3分钟",
        "5m": "5分钟",
        "15m": "15分钟",
        "30m": "30分钟",
        "1H": "1小时",
        "2H": "2小时",
        "4H": "4小时",
        "6H": "6小时",
        "12H": "12小时",
        "1D": "1天",
        "1W": "1周",
        "1M": "1月",
        "3M": "3月",
    }
    
    # 每个粒度最多可获取1440条
    MAX_LIMIT = 1440
    
    def __init__(self, sandbox: bool = True):
        """
        初始化指数K线数据管理器
        
        Args:
            sandbox: 是否使用模拟盘，True=模拟盘(1)，False=实盘(0)
        """
        if MarketData is None:
            raise ImportError("okx SDK未安装，请运行: pip install okx")
        
        self.sandbox = sandbox
        flag = "1" if sandbox else "0"
        
        # 初始化OKX行情数据API
        self.market_api = MarketData.MarketAPI(flag=flag)
        
        # 限速控制
        self.request_times: List[float] = []  # 请求时间戳列表（普通K线）
        self.history_request_times: List[float] = []  # 请求时间戳列表（历史K线）
        
        # 数据缓存
        self.candlesticks_cache: Dict[str, List[List]] = {}  # cache_key -> candlesticks
        self.cache_expire_time: Dict[str, datetime] = {}  # cache_key -> expire_time
        self.cache_duration = timedelta(seconds=30)  # 缓存30秒
        
        logger.info(f"指数K线数据管理器已初始化 (模式: {'模拟盘' if sandbox else '实盘'})")
    
    def _check_rate_limit(self, is_history: bool = False):
        """
        检查并执行限速
        
        Args:
            is_history: 是否为历史K线请求（限速更严格）
        """
        now = time.time()
        
        if is_history:
            # 历史K线限速：10次/2s (IP)
            self.history_request_times = [t for t in self.history_request_times 
                                         if now - t < self.HISTORY_RATE_LIMIT_WINDOW]
            
            if len(self.history_request_times) >= self.HISTORY_RATE_LIMIT_REQUESTS:
                oldest_request = self.history_request_times[0]
                wait_time = self.HISTORY_RATE_LIMIT_WINDOW - (now - oldest_request)
                
                if wait_time > 0:
                    logger.debug(f"触发历史K线限速，等待 {wait_time:.2f} 秒")
                    time.sleep(wait_time)
                    now = time.time()
                    self.history_request_times = [t for t in self.history_request_times 
                                                 if now - t < self.HISTORY_RATE_LIMIT_WINDOW]
            
            self.history_request_times.append(now)
        else:
            # 普通K线限速：20次/2s (IP)
            self.request_times = [t for t in self.request_times 
                                 if now - t < self.RATE_LIMIT_WINDOW]
            
            if len(self.request_times) >= self.RATE_LIMIT_REQUESTS:
                oldest_request = self.request_times[0]
                wait_time = self.RATE_LIMIT_WINDOW - (now - oldest_request)
                
                if wait_time > 0:
                    logger.debug(f"触发限速，等待 {wait_time:.2f} 秒")
                    time.sleep(wait_time)
                    now = time.time()
                    self.request_times = [t for t in self.request_times 
                                         if now - t < self.RATE_LIMIT_WINDOW]
            
            self.request_times.append(now)
    
    def get_index_candlesticks(self, inst_id: str, bar: str = "1H",
                              after: Optional[str] = None,
                              before: Optional[str] = None,
                              limit: Optional[int] = None,
                              use_cache: bool = True) -> List[List]:
        """
        获取指数K线数据
        
        Args:
            inst_id: 指数，如 BTC-USD
            bar: K线周期，如 1m/3m/5m/15m/30m/1H/2H/4H/6H/12H/1D/1W/1M/3M
            after: 请求此时间戳之前（更旧的数据）的分页内容，Unix时间戳（毫秒）
            before: 请求此时间戳之后（更新的数据）的分页内容，Unix时间戳（毫秒）
            limit: 返回结果的数量，最大1440，默认100
            use_cache: 是否使用缓存
            
        Returns:
            K线数据列表，每条K线包含：
            [0] ts: 开始时间（毫秒时间戳）
            [1] o: 开盘价格
            [2] h: 最高价格
            [3] l: 最低价格
            [4] c: 收盘价格
            [5] confirm: K线状态，0代表K线未完结，1代表K线已完结
        """
        # 验证K线周期
        if bar not in self.BAR_SIZES:
            raise ValueError(f"无效的K线周期: {bar}，支持的周期: {list(self.BAR_SIZES.keys())}")
        
        # 验证limit
        if limit is not None and (limit <= 0 or limit > self.MAX_LIMIT):
            raise ValueError(f"limit必须在1-{self.MAX_LIMIT}之间")
        
        # 生成缓存键
        cache_key = f"{inst_id}_{bar}_{after}_{before}_{limit}"
        
        # 检查缓存
        if use_cache and cache_key in self.candlesticks_cache:
            expire_time = self.cache_expire_time.get(cache_key)
            if expire_time and datetime.now() < expire_time:
                logger.debug(f"从缓存获取指数K线数据: {cache_key}")
                return self.candlesticks_cache[cache_key]
        
        # 执行限速检查
        self._check_rate_limit()
        
        try:
            # 构建请求参数
            params = {"instId": inst_id, "bar": bar}
            if after:
                params["after"] = after
            if before:
                params["before"] = before
            if limit:
                params["limit"] = str(limit)
            
            logger.info(f"获取指数K线数据: {inst_id} {self.BAR_SIZES[bar]} "
                       f"limit={limit or 100}")
            result = self.market_api.get_index_candlesticks(**params)
            
            # 检查响应
            if result.get("code") != "0":
                error_msg = result.get("msg", "未知错误")
                logger.error(f"获取指数K线数据失败: {error_msg}")
                raise RuntimeError(f"API错误: {error_msg}")
            
            # 提取数据
            data_list = result.get("data", [])
            logger.info(f"获取指数K线数据成功: {len(data_list)} 条")
            
            # 更新缓存
            if use_cache:
                self.candlesticks_cache[cache_key] = data_list
                self.cache_expire_time[cache_key] = datetime.now() + self.cache_duration
            
            return data_list
            
        except Exception as e:
            logger.error(f"获取指数K线数据异常: {e}")
            raise
    
    def get_latest_candlestick(self, inst_id: str, bar: str = "1H") -> Optional[Dict]:
        """
        获取最新的K线数据
        
        Args:
            inst_id: 指数
            bar: K线周期
            
        Returns:
            最新K线数据字典
        """
        candlesticks = self.get_index_candlesticks(inst_id, bar, limit=1)
        
        if candlesticks:
            return self._parse_candlestick(candlesticks[0])
        return None
    
    def get_history_candlesticks(self, inst_id: str, bar: str = "1D",
                                after: Optional[str] = None,
                                before: Optional[str] = None,
                                limit: Optional[int] = None,
                                use_cache: bool = True) -> List[List]:
        """
        获取指数历史K线数据（可获取最近几年的数据）
        
        Args:
            inst_id: 指数，如 BTC-USD
            bar: K线周期，如 1m/3m/5m/15m/30m/1H/2H/4H/6H/12H/1D/1W/1M/3M
            after: 请求此时间戳之前（更旧的数据）的分页内容，Unix时间戳（毫秒）
            before: 请求此时间戳之后（更新的数据）的分页内容，Unix时间戳（毫秒）
            limit: 返回结果的数量，最大100，默认100
            use_cache: 是否使用缓存
            
        Returns:
            K线数据列表，每条K线包含：
            [0] ts: 开始时间（毫秒时间戳）
            [1] o: 开盘价格
            [2] h: 最高价格
            [3] l: 最低价格
            [4] c: 收盘价格
            [5] confirm: K线状态，0代表K线未完结，1代表K线已完结
        """
        # 验证K线周期
        if bar not in self.BAR_SIZES:
            raise ValueError(f"无效的K线周期: {bar}，支持的周期: {list(self.BAR_SIZES.keys())}")
        
        # 验证limit（历史K线最大100）
        if limit is not None and (limit <= 0 or limit > 100):
            raise ValueError(f"历史K线limit必须在1-100之间")
        
        # 生成缓存键
        cache_key = f"history_{inst_id}_{bar}_{after}_{before}_{limit}"
        
        # 检查缓存
        if use_cache and cache_key in self.candlesticks_cache:
            expire_time = self.cache_expire_time.get(cache_key)
            if expire_time and datetime.now() < expire_time:
                logger.debug(f"从缓存获取指数历史K线数据: {cache_key}")
                return self.candlesticks_cache[cache_key]
        
        # 执行限速检查（历史K线限速更严格）
        self._check_rate_limit(is_history=True)
        
        try:
            # 构建请求参数
            params = {"instId": inst_id, "bar": bar}
            if after:
                params["after"] = after
            if before:
                params["before"] = before
            if limit:
                params["limit"] = str(limit)
            
            logger.info(f"获取指数历史K线数据: {inst_id} {self.BAR_SIZES[bar]} "
                       f"limit={limit or 100}")
            result = self.market_api.get_history_index_candlesticks(**params)
            
            # 检查响应
            if result.get("code") != "0":
                error_msg = result.get("msg", "未知错误")
                logger.error(f"获取指数历史K线数据失败: {error_msg}")
                raise RuntimeError(f"API错误: {error_msg}")
            
            # 提取数据
            data_list = result.get("data", [])
            logger.info(f"获取指数历史K线数据成功: {len(data_list)} 条")
            
            # 更新缓存（历史数据缓存时间更长）
            if use_cache:
                self.candlesticks_cache[cache_key] = data_list
                self.cache_expire_time[cache_key] = datetime.now() + timedelta(hours=1)
            
            return data_list
            
        except Exception as e:
            logger.error(f"获取指数历史K线数据异常: {e}")
            raise
    
    def get_candlesticks_range(self, inst_id: str, bar: str,
                              start_time: datetime, end_time: datetime,
                              use_history_api: bool = False) -> List[Dict]:
        """
        获取指定时间范围的K线数据
        
        Args:
            inst_id: 指数
            bar: K线周期
            start_time: 开始时间
            end_time: 结束时间
            use_history_api: 是否使用历史K线API（可获取更长时间的数据，但限速更严格）
            
        Returns:
            K线数据列表
        """
        # 转换为毫秒时间戳
        before_ts = str(int(end_time.timestamp() * 1000))
        after_ts = str(int(start_time.timestamp() * 1000))
        
        all_candlesticks = []
        current_before = before_ts
        
        # 根据时间跨度自动选择API
        time_diff = end_time - start_time
        if not use_history_api and time_diff > timedelta(days=60):
            logger.info(f"时间跨度超过60天，自动使用历史K线API")
            use_history_api = True
        
        while True:
            # 获取一批数据
            if use_history_api:
                candlesticks = self.get_history_candlesticks(
                    inst_id=inst_id,
                    bar=bar,
                    before=current_before,
                    limit=100,  # 历史API最大100
                    use_cache=False
                )
            else:
                candlesticks = self.get_index_candlesticks(
                    inst_id=inst_id,
                    bar=bar,
                    before=current_before,
                    limit=self.MAX_LIMIT,
                    use_cache=False
                )
            
            if not candlesticks:
                break
            
            # 过滤时间范围
            for candle in candlesticks:
                ts = int(candle[0])
                if int(after_ts) <= ts <= int(before_ts):
                    all_candlesticks.append(self._parse_candlestick(candle))
            
            # 检查是否已经获取到开始时间之前的数据
            oldest_ts = int(candlesticks[-1][0])
            if oldest_ts <= int(after_ts):
                break
            
            # 更新分页参数
            current_before = str(oldest_ts)
            
            # 避免无限循环
            if len(all_candlesticks) > 10000:
                logger.warning("获取的K线数据超过10000条，停止获取")
                break
        
        # 按时间排序（从旧到新）
        all_candlesticks.sort(key=lambda x: x["timestamp_ms"])
        
        logger.info(f"获取时间范围K线数据: {len(all_candlesticks)} 条")
        return all_candlesticks
    
    def _parse_candlestick(self, candle: List) -> Dict:
        """
        解析K线数据
        
        Args:
            candle: 原始K线数据
            
        Returns:
            解析后的字典
        """
        ts = int(candle[0])
        return {
            "timestamp_ms": ts,
            "timestamp": datetime.fromtimestamp(ts / 1000).strftime("%Y-%m-%d %H:%M:%S"),
            "open": float(candle[1]),
            "high": float(candle[2]),
            "low": float(candle[3]),
            "close": float(candle[4]),
            "confirm": int(candle[5]),
            "is_complete": int(candle[5]) == 1,
        }
    
    def calculate_indicators(self, candlesticks: List[Dict]) -> Dict:
        """
        计算K线指标
        
        Args:
            candlesticks: K线数据列表
            
        Returns:
            指标字典
        """
        if not candlesticks:
            return {}
        
        # 提取价格数据
        closes = [c["close"] for c in candlesticks]
        highs = [c["high"] for c in candlesticks]
        lows = [c["low"] for c in candlesticks]
        
        # 计算基本指标
        current_price = closes[-1]
        highest = max(highs)
        lowest = min(lows)
        price_range = highest - lowest
        
        # 计算涨跌幅
        if len(closes) > 1:
            first_price = closes[0]
            change = current_price - first_price
            change_pct = (change / first_price * 100) if first_price > 0 else 0
        else:
            change = 0
            change_pct = 0
        
        # 计算简单移动平均线
        ma_periods = [5, 10, 20, 50]
        mas = {}
        for period in ma_periods:
            if len(closes) >= period:
                mas[f"ma{period}"] = sum(closes[-period:]) / period
        
        return {
            "current_price": current_price,
            "highest": highest,
            "lowest": lowest,
            "price_range": price_range,
            "change": change,
            "change_pct": change_pct,
            "candle_count": len(candlesticks),
            **mas
        }
    
    def get_candlesticks_with_indicators(self, inst_id: str, bar: str = "1H",
                                        limit: int = 100) -> Tuple[List[Dict], Dict]:
        """
        获取K线数据并计算指标
        
        Args:
            inst_id: 指数
            bar: K线周期
            limit: 数量
            
        Returns:
            (K线数据列表, 指标字典)
        """
        # 获取K线数据
        raw_candlesticks = self.get_index_candlesticks(inst_id, bar, limit=limit)
        candlesticks = [self._parse_candlestick(c) for c in raw_candlesticks]
        
        # 计算指标
        indicators = self.calculate_indicators(candlesticks)
        
        return candlesticks, indicators
    
    def compare_multiple_indexes(self, inst_ids: List[str], bar: str = "1H",
                                limit: int = 100) -> Dict[str, Dict]:
        """
        比较多个指数的K线数据
        
        Args:
            inst_ids: 指数列表
            bar: K线周期
            limit: 数量
            
        Returns:
            字典，key为指数，value为指标数据
        """
        results = {}
        
        for inst_id in inst_ids:
            try:
                _, indicators = self.get_candlesticks_with_indicators(
                    inst_id, bar, limit
                )
                results[inst_id] = indicators
            except Exception as e:
                logger.error(f"获取 {inst_id} K线数据失败: {e}")
                results[inst_id] = None
        
        return results
    
    def clear_cache(self, inst_id: Optional[str] = None):
        """
        清除缓存
        
        Args:
            inst_id: 指数，如果为None则清除所有缓存
        """
        if inst_id:
            # 清除所有包含该inst_id的缓存
            keys_to_remove = [key for key in self.candlesticks_cache.keys() 
                            if key.startswith(f"{inst_id}_")]
            for key in keys_to_remove:
                del self.candlesticks_cache[key]
                if key in self.cache_expire_time:
                    del self.cache_expire_time[key]
            logger.info(f"已清除 {inst_id} 的缓存")
        else:
            self.candlesticks_cache.clear()
            self.cache_expire_time.clear()
            logger.info("已清除所有指数K线缓存")
    
    def get_cache_stats(self) -> Dict:
        """
        获取缓存统计信息
        
        Returns:
            缓存统计字典
        """
        total_cached = len(self.candlesticks_cache)
        valid_cached = sum(1 for key in self.candlesticks_cache.keys()
                          if key in self.cache_expire_time and
                          datetime.now() < self.cache_expire_time[key])
        
        return {
            "total_cached": total_cached,
            "valid_cached": valid_cached,
            "expired_cached": total_cached - valid_cached,
            "cache_duration_seconds": self.cache_duration.total_seconds(),
        }


# 便捷函数
def get_index_candlesticks(inst_id: str, bar: str = "1H", limit: int = 100,
                          sandbox: bool = True) -> List[Dict]:
    """
    便捷函数：获取指数K线数据
    
    Args:
        inst_id: 指数
        bar: K线周期
        limit: 数量
        sandbox: 是否使用模拟盘
        
    Returns:
        K线数据列表
    """
    manager = IndexCandlesticksManager(sandbox=sandbox)
    raw_candlesticks = manager.get_index_candlesticks(inst_id, bar, limit=limit)
    return [manager._parse_candlestick(c) for c in raw_candlesticks]


def get_latest_candlestick(inst_id: str, bar: str = "1H",
                          sandbox: bool = True) -> Optional[Dict]:
    """
    便捷函数：获取最新K线
    
    Args:
        inst_id: 指数
        bar: K线周期
        sandbox: 是否使用模拟盘
        
    Returns:
        最新K线数据
    """
    manager = IndexCandlesticksManager(sandbox=sandbox)
    return manager.get_latest_candlestick(inst_id, bar)


def get_candlesticks_with_indicators(inst_id: str, bar: str = "1H",
                                    limit: int = 100, sandbox: bool = True) -> Tuple[List[Dict], Dict]:
    """
    便捷函数：获取K线数据并计算指标
    
    Args:
        inst_id: 指数
        bar: K线周期
        limit: 数量
        sandbox: 是否使用模拟盘
        
    Returns:
        (K线数据列表, 指标字典)
    """
    manager = IndexCandlesticksManager(sandbox=sandbox)
    return manager.get_candlesticks_with_indicators(inst_id, bar, limit)


def get_history_candlesticks(inst_id: str, bar: str = "1D", limit: int = 100,
                             sandbox: bool = True) -> List[Dict]:
    """
    便捷函数：获取指数历史K线数据
    
    Args:
        inst_id: 指数
        bar: K线周期
        limit: 数量（最大100）
        sandbox: 是否使用模拟盘
        
    Returns:
        K线数据列表
    """
    manager = IndexCandlesticksManager(sandbox=sandbox)
    raw_candlesticks = manager.get_history_candlesticks(inst_id, bar, limit=limit)
    return [manager._parse_candlestick(c) for c in raw_candlesticks]


def get_candlesticks_range(inst_id: str, bar: str, start_time: datetime,
                          end_time: datetime, sandbox: bool = True) -> List[Dict]:
    """
    便捷函数：获取指定时间范围的K线数据
    
    Args:
        inst_id: 指数
        bar: K线周期
        start_time: 开始时间
        end_time: 结束时间
        sandbox: 是否使用模拟盘
        
    Returns:
        K线数据列表
    """
    manager = IndexCandlesticksManager(sandbox=sandbox)
    return manager.get_candlesticks_range(inst_id, bar, start_time, end_time)
